#!D:/software/Miniconda3/envs/fast-lane/python
#coding=utf-8

import cv2
from model.model import parsingNet
import torch
import scipy.special
import numpy as np
import torchvision.transforms as transforms
from PIL import Image
import time

class simulation_opt:
    def __init__(self,
            weights = None, num_lanes=None, device = 'cpu',
            # 图森数据格式
            dataset = 'Tusimple',
            # resNet18
            backbone='18',
            # 网格点100
            griding_num = 100,
           ):

        self.backbone = backbone
        self.dataset = dataset
        self.griding_num = griding_num
        self.test_model = weights
        self.num_lanes = num_lanes
        self.device = device
        self.tusimple_row_anchor = [ 64,  68,  72,  76,  80,  84,  88,  92,  96, 100, 104, 108, 112,
            116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164,
            168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216,
            220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 260, 264, 268,
            272, 276, 280, 284]

class lane_det_api:
    def __init__(self,weights, num_lanes, device):
        self.opt = simulation_opt(weights = weights,num_lanes = num_lanes, device = device )
        # 加载参数
        torch.backends.cudnn.benchmark = True # 那么cuDNN使用的非确定性算法就会自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题
        print('start testing...')
        if (self.opt.dataset == 'Tusimple'):
            self.cls_num_per_lane = 56
            self.row_anchor = self.opt.tusimple_row_anchor
        else:
            raise NotImplementedError

        self.net = parsingNet(pretrained=False, backbone=self.opt.backbone, cls_dim=(self.opt.griding_num + 1, self.cls_num_per_lane, self.opt.num_lanes),
                         use_aux=False).to(self.opt.device)  # we dont need auxiliary segmentation in testing

        state_dict = torch.load(self.opt.test_model, map_location='cpu')['model']
        compatible_state_dict = {}

        for k, v in state_dict.items():
            if 'module.' in k:
                compatible_state_dict[k[7:]] = v
            else:
                compatible_state_dict[k] = v

        self.net.load_state_dict(compatible_state_dict, strict=False)
        self.net.eval()

        # 图像变换
        self.img_transforms = transforms.Compose([
            transforms.Resize((288, 800)),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])

    def lane_det(self, frame):

        img_w = 1280
        img_h = 720
        frame = cv2.resize(frame, (img_w, img_h))
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img_ = Image.fromarray(img)  # 实现array到image的转换

        imgs = self.img_transforms(img_)
        imgs = imgs.unsqueeze(0)  # 起到升维的作用

        imgs = imgs.to(self.opt.device)

        t1 = time.time()

        with torch.no_grad():
            out = self.net(imgs)

        t2 = time.time()

        col_sample = np.linspace(0, 800 - 1, self.opt.griding_num)
        col_sample_w = col_sample[1] - col_sample[0]

        out_j = out[0].data.cpu().numpy()
        out_j = out_j[:, ::-1, :]
        prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
        idx = np.arange(self.opt.griding_num) + 1
        idx = idx.reshape(-1, 1, 1)
        loc = np.sum(prob * idx, axis=0)
        out_j = np.argmax(out_j, axis=0)
        loc[out_j == self.opt.griding_num] = 0
        out_j = loc

        t3 = time.time()
        print('Done.({:.2}s)'.format(t2 - t1))
        points = 0
        for i in range(out_j.shape[1]):
            if np.sum(out_j[:, i] != 0) > 2:
                # print("--------------")
                point = []
                for k in range(out_j.shape[0]):
                    if out_j[k, i] > 0:
                        ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1,
                               int(img_h * (self.row_anchor[self.cls_num_per_lane - 1 - k] / 288)) - 1)  # 恢复尺度
                        # ----------------统计了画线段----------------------------
                        point.append(ppp)

                for i, p in enumerate(point):
                    cv2.circle(frame, p, 3, (0, 255, 0), -1)
                    if (i < len(point) - 1):  # 最后一个点不需要画线段
                        cv2.line(frame, p, point[i + 1], (255, 0, 0), 2)
                points = point
        return frame,points

if __name__ == "__main__":
    # 读取数据
    cap = cv2.VideoCapture("3.mp4")
    weights_path = r"D:\code\Fast-Lane-Detection\logs\1027_0936_lr_4e-04_b_32\ep280.pth"

    # 修改权重路径，车道线数，是否选用显卡device = 0 or 'cpu'
    DetectLane= lane_det_api(weights = weights_path, num_lanes = 2, device = 0)
    while True:
        rec, img = cap.read()
        if rec == 0:
            print("Image source is error...")
            break
        # 返回图片和点集
        img_result, points = DetectLane.lane_det(img)
        cv2.imshow("result", img_result)
        if(cv2.waitKey(1) == (ord('q'))):
            break


